
from typing import List
import os
import datetime
from PIL import Image
import numpy as np
# import folder_paths
import tensorflow as tf
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
from alioffline_utils_img import *
from alioffline_utils_mask import png_to_mask,png_add_bg

# comfy_path = os.path.dirname(folder_paths.__file__)
current_directory = os.path.dirname(__file__)
comfy_path =  os.path.abspath(os.path.join(current_directory, '..', '..'))
models_path = os.path.join(comfy_path, "models","modelscope")
# model_path = os.path.join(models_path, "damo","cv_unet_universal-matting")

class AliOffline_Seg_Obj:
   
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
               {   
                "image": ('IMAGE', {}),
                "image_type":   (["whiteBK","blackBK", "transparent"],{"default":"transparent"} ),   
                "mask_type":   (["whiteBK","blackBK", "transparent"],{"default":"transparent"} ),   
              }
        }

    RETURN_TYPES = ("IMAGE","IMAGE",)
    RETURN_NAMES = ("image","mask",)
    OUTPUT_NODE = True
    FUNCTION = "sample"
    CATEGORY = "CXH"

    def sample(self,image,image_type="transparent",mask_type="transparent"):
        os.environ['MODELSCOPE_CACHE']=models_path
        #res_images = []
        #res_masks = []
        tf.compat.v1.disable_eager_execution()
        sess = tf.compat.v1.Session()
        # model = pipeline(
        #     Tasks.portrait_matting, model="damo/cv_unet_image-matting"
        # )
        model = pipeline(
            Tasks.universal_matting, model="damo/cv_unet_universal-matting"
        )
        #for item in image:
        result = model(tensor2pil(image))
        ndarray = cv2.cvtColor(result[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGBA)
        out_image = array2image(ndarray)
        out_tensor = pil2tensor(out_image)
        if image_type == "whiteBK":
          img_with_bg = png_add_bg(out_image,255,255,255)
          out_tensor = pil2tensor(img_with_bg)
        #  res_images.extend(out_tensor)
        elif mask_type == "blackBK":
          img_with_bg = png_add_bg(out_image,0,0,0)
          out_tensor = pil2tensor(img_with_bg)
        #  res_images.extend(out_tensor)
        #else:
        #  out_tensor = pil2tensor(out_image)
        #  res_images.extend(out_tensor)

        out_mask = pil2tensor(png_to_mask(out_image,False,0))
        if mask_type == "whiteBK":
          out_mask = pil2tensor(png_to_mask(out_image,False,255))
        #  res_masks.extend(out_mask)
        elif mask_type == "blackBK":
          out_mask = pil2tensor(png_to_mask(out_image,True,255))
        #  res_masks.extend(out_mask)
        #else:
        #  out_mask = pil2tensor(png_to_mask(out_image,False,0))
        #  res_masks.extend(out_mask)

        sess.close()
        return (out_tensor,out_mask)